#include "acl/acl.h"
#include "../../include/custom_type.h"
#include "../data_utils.h"
#include "../run_main.h"
#include "./LLMsGEMM_batch_QKTV_host.h"
#include <cstdio>
#include <string>



int main ( int argc, char ** argv ) {

    // 获取输入 layoutA layoutB zeroPaddingM zeroPaddingN zeroPaddingK batchCount
    layoutType layoutA = ( std::stoi( argv[1] ) == 0 ? RowMajor : ColumnMajor );
    layoutType layoutB = ( std::stoi( argv[2] ) == 0 ? RowMajor : ColumnMajor );
    int64_t zeroPaddingM = std::stoi( argv[3] );
    int64_t zeroPaddingN = std::stoi( argv[4] );
    int64_t zeroPaddingK = std::stoi( argv[5] );
    int64_t batchCount = std::stoi( argv[6] );  
    
    // 打印输出 layoutA layoutB zeroPaddingM zeroPaddingN zeroPaddingK batchCount
    printf("\n\nlayoutA: %s, layoutB: %s, "
            "zeroPaddingM: %ld, zeroPaddingN: %ld, zeroPaddingK: %ld, " 
            "batchCount: %ld. \n\n", 
            layoutA == 0 ? "RowMajor" : "ColumnMajor", layoutB == 0 ? "RowMajor" : "ColumnMajor", 
            zeroPaddingM, zeroPaddingN, zeroPaddingK, 
            batchCount);

    // acl初始化
    int device = 4;
    const char *aclConfigPath = "../acl.json";
    ACL_CHECK(aclInit(/*aclConfigPath*/nullptr));
    ACL_CHECK(aclrtSetDevice(device));
    aclrtStream stream;
    ACL_CHECK(aclrtCreateStream(&stream));

    std::string src="./data/";
    std::string path_A = src + "A.bin";
    std::string path_B = src + "B.bin";
    std::string path_C = src + "C.bin";
    std::string path_alpha = src + "alpha.bin";
    std::string path_beta  = src + "beta.bin";
    std::string path_validM = src + "validM.bin";
    std::string path_expect_result = src + "expect_result.bin";

    half* pre_output = nullptr;
    half* cur_output = nullptr;

    run_main(
        LLMsGEMM_batch_QKTV_host,
        stream, 
        layoutA, 
        layoutB, 
        zeroPaddingM, 
        zeroPaddingN, 
        zeroPaddingK, 
        batchCount, 
        path_A, 
        path_B, 
        path_C, 
        path_alpha, 
        path_beta, 
        path_validM, 
        path_expect_result, 
        pre_output, 
        cur_output, 
        2, 
        0
    );

    ACL_CHECK(aclrtFree(cur_output));

    // 反初始化
    ACL_CHECK(aclrtDestroyStream(stream));
    ACL_CHECK(aclrtResetDevice(device));
    ACL_CHECK(aclFinalize());

    return 0;
}